import logging
import math
import gym
from gym import spaces
from gym.utils import seeding
import numpy as np
import random
import os
import torch
import torch.nn as nn

from Comparison_Algorithm.Highest_Bitrate import Env_Alg as HB

from queue import Queue

# DIR = "data_test"
DIR = "data_train"

logger = logging.getLogger(__name__)

class DataSet():
    def __init__(self):
        # 二维向量，横轴是分片，纵轴是时隙
        self.bitratelevel = []
        self.bitrateInfo = []
        self.buffer = []
        self.FovRecord = []
        self.rebuffer_time = []
        self.throughput = []
        self.pause_total = []
        self.name = ""



    def Fetch_data_from_file(self,filepath,filename):
        pos = []
        file = filepath+filename
        with open(file, 'r') as file_to_read:
            while True:
                lines = file_to_read.readline()  # 整行读取数据
                if not lines:
                    break
                p_tmp = [float(i) for i in lines.split(',')]  # 将整行数据分割处理，如果分割符是空格，括号里就不用传入参数，如果是逗号， 则传入‘，'字符。
                pos.append(p_tmp)  # 添加新读取的数据
            pos = np.array(pos)
        if filename == "bitrateLevel.txt":
            self.bitratelevel = pos
        elif filename == "bitrateInfo.txt":
            self.bitrateInfo = pos
        elif filename == "throughput.txt":
            self.throughput = pos
        elif filename == "buffer.txt":
            self.buffer = pos
        elif filename == "FovRecord.txt":
            self.FovRecord = pos
        elif filename == "pause_total.txt":
            self.pause_total = pos
        elif filename == "rebuffer_time.txt":
            self.rebuffer_time = pos




class TiledVideo(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second' : 50
    }

    def __init__(self):

        self.video = np.array([0,1,2,3,4,5]) # 2*2 tiled  不懂 video如何
        self.action_space = spaces.Discrete(4*4*4*4*4*4) # 4^6 每一个tile可以选择4种码率
        #self.action_space = spaces.MultiDiscrete(4 * np.ones((4,)))
        # 前六个FoV，中6个buffer，后6个吞吐量
        self.observation_space = spaces.Discrete(7) # 6 FoV+1 buffer
        # self.observation_space.n = 14
        #self.observation_space = spaces.MultiDiscrete(4 * np.ones((4,)))
        self.rng = np.random.RandomState(1234) # 1234是随机数种子
        self.state = [0,0,0,0,0,0,0] # * 2 #
        self.ep_steps = 0 # 
        self.video_slot = 0
        self.record_bitratelevel = []
        self.dirs = [] # 存文件数据文件名的
        # self.tile_adjacent_relation = {0:[1,2],
        #                                1:[0,3],
        #                                2:[0,3],
        #                                3:[1,2]}
        self.tile_adjacent_relation = {0:[2,3,4,5],
                                       1:[2,3,4,5],
                                       2:[0,1,4,5],
                                       3:[0,1,4,5],
                                       4:[0,1,2,3],
                                       5:[0,1,2,3]}  # tile与其他tile的关联程度，是自行设计的
        self.buffer = [0,0,0,0,0,0]  # 每一个tile有一个buffer
        self.last_bitrate = [0,0,0,0,0,0] # [[0,0,0,0]] # 4个tile 1个时隙  上一时刻的比特率
        #self.last_throughout = [] 
        self.rebuffer_event = [] # rebuffer事件

        # 读取数据集信息
        # self.dataset = DataSet() # 数据存储类
        self.dataset_all = [DataSet() for i in range(len(os.listdir("./"+DIR)))]
        self.dataset = DataSet()
        self.load_data_all()

        # Comparison Algorithm
        self.HB_Comparison = HB()
        self.DP_Comparison = HB()

        # record
        self.last_bitrate = [] # 上一个帧的比特率
        self.last_throughout = []

    def load_data_all(self):
        # self.dirs = os.listdir("D:\Project_Python\ABR\option-critic-pytorch-master\data_train")
        self.dirs = os.listdir("./"+DIR)

        # 全读进内存
        for i in range(len(self.dirs)):
            dir_path = DIR + "/" + self.dirs[i] + "/"
            self.dataset_all[i].name = dir_path
            self.dataset_all[i].Fetch_data_from_file(dir_path, "bitrateLevel.txt")
            self.dataset_all[i].Fetch_data_from_file(dir_path, "bitrateInfo.txt")
            self.dataset_all[i].Fetch_data_from_file(dir_path, "throughput.txt")
            self.dataset_all[i].Fetch_data_from_file(dir_path, "buffer.txt")
            self.dataset_all[i].Fetch_data_from_file(dir_path, "FovRecord.txt")
            self.dataset_all[i].Fetch_data_from_file(dir_path, "pause_total.txt")
            self.dataset_all[i].Fetch_data_from_file(dir_path, "rebuffer_time.txt")

    def load_data_one(self):
        # dir_path = DIR + "/" + random.choice(self.dirs) + "/"
        self.dataset = random.choice(self.dataset_all)
        print(self.dataset.name)

    def seed(self, seed=None):
        return self._seed(seed)

    def _seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def reset(self):
        self.state = [0, 0, 0, 0, 0, 0, 0]# *2
        self.buffer = [0,0,0,0,0,0]
        # self.last_bitrate = []
        self.last_bitrate = [] # [[0,0,0,0]] # 4个tile 1个时隙
        self.last_throughout = []
        self.rebuffer_event = []
        self.ep_steps = 0 # 没用上啊，main_ABR的程序里自己有记录
        self.record_bitratelevel = []  # 归零更新
        self.load_data_one()  # 在这里读文件
        # self.video_slot = 0  ## reset有两种情况，一个是视频播完，另一个是步长太大，所以不能在这个reset里归零video_slot
        return self.state

    def ten_to_4(self, x): # 10转4进制
        output = ""
        ret = x // 4
        leave = x % 4
        if ret == 0:
            return str(leave)
        xxx = self.ten_to_4(ret)
        output = xxx + str(leave)
        return output

    def Dimensionality_reduction(self,FoV_4_6):
        FoV_6 = [0,0,0,0,0,0]
        for i in range(6):
            FoV_6[i] = 0.5*FoV_4_6[0][i] + FoV_4_6[1][i] + FoV_4_6[2][i]

        return FoV_6

    def step(self, action, current_option): # 将已经判断出的action current_option送入，获取reward？

        # action = action.argmax()
        rebuffer_record = 0
        FoV = self.dataset.FovRecord
        throughput = self.dataset.throughput
        done = False
        throughput = throughput/8000000

        throughput_mean = np.mean(throughput[self.video_slot])
        throughput_sum = np.sum(throughput[self.video_slot])

        # state_rate_mapping = {0:777451,
        #                       1:2769991,
        #                       2:6430841,
        #                       3:116736619}
        # 太大，导致状态空间太大，更新后的数值太大
        # state_rate_mapping = {0: 100000000,
        #                       1: 300000000,
        #                       2: 500000000,
        #                       3: 700000000} # 单位bps
        # state_rate_mapping = {0: 12.5,
        #                       1: 37.5,
        #                       2: 62.5,
        #                       3: 87.5} # 单位MBps
        state_rate_mapping = {0: 8,
                              1: 24,
                              2: 36,
                              3: 48} # 单位MBps

        FoV_Adjacent = {0:[2,3,4,5],
                        1:[2,3,4,5],
                        2:[0,1,4,5],
                        3:[0,1,4,5],
                        4:[0,1,2,3],
                        5:[0,1,2,3]}

        action = [action[i].detach().item() for i in range(6)]
        self.record_bitratelevel.append(self.state[0:6]) # 记录之前的码率选择

        rebuffer_event = 0
        rebuffer_flag = 0
        for i in range(6):
            if self.buffer[i] < 1:
                self.buffer[i] += throughput_mean / state_rate_mapping[action[i]]
                if self.buffer[i] < 1:
                    rebuffer_event += 1-self.buffer[i]
                    rebuffer_flag  = 1
                    # self.buffer[i] = 1
            else:
                self.buffer[i] = 1 + throughput_mean / state_rate_mapping[action[i]]
        if not rebuffer_flag:
            self.buffer = (np.array(self.buffer) - 1).tolist()   # 模拟buffer过程
        # 现在的逻辑是第一个FoV为1的当成主视角，也就是最高优先级 主视角0.8 主视角相邻的tile中 有1的给0.2，0的给0.1
        artificial_range_1s = {}
        range_temp = [0,1,2,3,4,5]
        main_perspective = 0
        main_near = []
        artificial_range_2s = {}

        artificial_main = 1
        artificial_main_near = 0.2
        artificial_not_main = 0.1
        artificial_duimian = 0
        for i in range_temp:
            if FoV[self.video_slot + 1][i].astype(int) == 1:
                # count = list(FoV[self.video_slot + 1].astype(int)).count(1)
                artificial_range_1s[i] = artificial_main  # 主视野内权重为0.8
                range_temp.remove(i)
                main_perspective = i
                for j in FoV_Adjacent[i]:
                    range_temp.remove(j)
                    if FoV[self.video_slot + 1][j].astype(int) == 1:
                        artificial_range_1s[j] = artificial_main_near# 0.5
                        main_near.append(j)
                    else:
                        artificial_range_1s[j] = artificial_not_main# 0.2

            artificial_range_1s[range_temp[0]] = artificial_duimian ## ???

        bitrate_level_total = 0
        # 乘哪个FoV这个判断直接在state中考量，，等一下，有点问题
        for i in range(6):  # 6个面
            bitrate_level_total += action[i] * artificial_range_1s[i]

        # 帧内平滑度
        # 考虑FoV的东西，这个平滑度考虑主视角和其他四片的差异
        smoothness_within_frame = 0
        for i in FoV_Adjacent[main_perspective]:
            if i in main_near:
                smoothness_within_frame += abs(action[main_perspective] - action[i])
            else:
                smoothness_within_frame += abs(0.5 * (action[main_perspective] - action[i]))

        bitrate = 0
        for i in range(6):
            bitrate += action[i]
        # 1.8-0.3-0.5
        alpha = 3 #.5
        beta = 1.5
        gamma = 4
        # bitrate_factor=0.4 #3 # 0.8会导致全选0
        # reward = alpha * bitrate_level_total - (bitrate_factor* bitrate) - beta * smoothness_within_frame - gamma*rebuffer_event
        # reward = alpha * bitrate_level_total - bitrate_factor * bitrate - gamma * rebuffer_event
        # reward = alpha * bitrate_level_total - gamma * rebuffer_event
        reward = alpha * bitrate_level_total - beta * smoothness_within_frame - gamma * rebuffer_event

        reward_HB = self.HB_Comparison.Highest_Bitrate_Selection(FoV, throughput_mean,alpha,beta,gamma,
                                                                 artificial_main,artificial_main_near,artificial_not_main,artificial_duimian)#,bitrate_factor)
        reward_DP = self.DP_Comparison.Dynamic_Profit(FoV[self.video_slot + 1],FoV[self.video_slot + 2],throughput_mean,
                                                      alpha,beta,gamma,artificial_main,artificial_main_near,artificial_not_main,artificial_duimian)
        #with open(DIR+"_result.txt","a") as f:
        #    f.write("FoV2: "+str(FoV[self.video_slot+1])+",  因FoV做的决策: " + str(action)+ ",   option: " + str(current_option) +  "\n")

        self.video_slot += 1

        if self.video_slot >= len(self.dataset.throughput)-2:
            done = True
            self.video_slot = 0


        FoV_state = FoV[self.video_slot+1]
        self.state[0:6] = FoV_state # FoV应该给未来时刻的
        self.state[6] = throughput_mean # throughput_sum
        #self.state[6:12] = self.buffer # buffer应该给当前时刻的
        # self.state[12:] = throughput[self.video_slot-1] # 吞吐量应该给前一时刻的

        # for i in range(0,6):
        #     self.state[i] = FoV_state[i]
        latency = 0;data_amount = 0;##这俩从系统上测吧？
        _ = [bitrate_level_total,smoothness_within_frame,rebuffer_event, latency, data_amount,
             reward_HB,reward_DP]# reward_HB
        return self.state, reward, done, _

if __name__=="__main__":
    env = TiledVideo()
    env.seed(3)
    print("pause")
